#pragma once
#include <windows.h>
#include <iostream>
using namespace std;

class ScreenGrab {
    bool _debug;
public:
    void setDebug(bool debug = false) {
        _debug = debug;
    }
    bool screenCapture(const wchar_t* filename, int x = 0, int y = 0, int w = 1536, int h = 864) {
        // system("chcp 65001");
        //1.获取主屏幕的显示设备的上下文
        HDC hdcscreen = GetDC(NULL);
        int screenWidth = GetDeviceCaps(hdcscreen, HORZRES);
        int screenHeight = GetDeviceCaps(hdcscreen, VERTRES);
        int bitsPerPixel = GetDeviceCaps(hdcscreen, BITSPIXEL);
        if (_debug) {
            cout << "主屏幕的分辨率: " << screenWidth << "x" << screenHeight << endl;
            cout << "位深: " << bitsPerPixel << endl;
        }
        HDC hMemDC = CreateCompatibleDC(hdcscreen);
        if (hMemDC == NULL) {
            DWORD err = GetLastError();
            std::cerr << "创建位图失败! 错误码: " << err << std::endl;
            DeleteDC(hdcscreen);
            ReleaseDC(NULL, hdcscreen);
            return -1;
        }
        HBITMAP hBitMap = CreateCompatibleBitmap(hdcscreen, 1536, 864);
        if (hBitMap == NULL) {
            DWORD err = GetLastError();
            // std::cout << "hFile = " << hBitMap << ", INVALID_HANDLE_VALUE = " << INVALID_HANDLE_VALUE << std::endl;
            std::cerr << "创建位图失败! 错误码: " << err << std::endl;
            DeleteDC(hdcscreen);
            ReleaseDC(NULL, hdcscreen);
            return -1;
        }
        SelectObject(hMemDC, hBitMap);
        BitBlt(hMemDC, 0, 0, w, h, hdcscreen, x, y, SRCCOPY);
        BITMAPFILEHEADER bmpFileHeader = {0};
        BITMAPINFOHEADER bmpInfoHeader = {0};
        bmpFileHeader.bfType = 0x4D42;
        bmpFileHeader.bfSize = sizeof(BITMAPFILEHEADER) + 
                            sizeof(BITMAPINFOHEADER) + 
                            (w * h * bitsPerPixel / 8);
        bmpFileHeader.bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);

        bmpInfoHeader.biSize = sizeof(BITMAPINFOHEADER);
        bmpInfoHeader.biHeight = -h;
        bmpInfoHeader.biWidth = w;
        bmpInfoHeader.biPlanes = 1;
        bmpInfoHeader.biBitCount = bitsPerPixel;
        bmpInfoHeader.biCompression = BI_RGB;

        HANDLE hFile = CreateFileW(filename, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
        DWORD err = GetLastError();
        // std::cerr << "hFile = " << hFile << ", INVALID_HANDLE_VALUE = " << INVALID_HANDLE_VALUE << std::endl;
        if (hFile == INVALID_HANDLE_VALUE) {
            std::cerr << "创建bitmap图片失败! 错误码: " << err << std::endl;
            DeleteObject(hBitMap);
            DeleteDC(hMemDC);
            ReleaseDC(NULL, hdcscreen);
            return -1;
        }

        DWORD bytesWritten = 0;
        WriteFile(hFile, &bmpFileHeader, sizeof(bmpFileHeader), &bytesWritten, NULL);
        WriteFile(hFile, &bmpInfoHeader, sizeof(bmpInfoHeader), &bytesWritten, NULL);
        DWORD pixelDataSize = w * h * bitsPerPixel / 8;
        BYTE* pixelBuffer = new BYTE[pixelDataSize];
        GetDIBits(hMemDC, hBitMap, 0, h, pixelBuffer, (BITMAPINFO*)&bmpInfoHeader, DIB_RGB_COLORS);
        WriteFile(hFile, pixelBuffer, pixelDataSize, &bytesWritten, NULL);

        delete [] pixelBuffer;
        CloseHandle(hFile);
        DeleteObject(hBitMap);
        DeleteDC(hMemDC);
        ReleaseDC(NULL, hdcscreen);
        if (_debug) {
            std::wcout << L"BMP截图保存成功: " << filename << std::endl;
        }
        return true;
    }
};